﻿/*
 * DHKex.java
 *
 * Created on 2007年8月21日, 下午 3:14
 *
 */

package org.uranus.core.ssh.v2;

import java.util.Random;
import org.uranus.core.ssh.SSHConnectionException;

/**
 * Copyright 2007 Hsiang-Yu Wang<br>
 * This file is distributed under the terms of the GNU General Public License<br>
 *<br>
 *<pre>
 * --------------------License notice--------------------
 * GNU GENERAL PUBLIC LICENSE
 * Version 3, 29 June 2007
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not,
 * visit: http://www.gnu.org/licenses/gpl.txt
 * --------------------License notice--------------------</pre>
 *
 * @author Hsiang-Yu Wang
 */
public class DHKex{
    private BigInteger p;
    private BigInteger g;
    private BigInteger x;
    private BigInteger y;
    
    private byte[] e;
    
    private byte[] K_S;
    
    private SshPacketV2 pn;
    private byte[] sign_of_H;
    private byte[] K;
    
    private byte[] V_S;
    private byte[] H;
    private byte[] V_C;
    
    private byte[] I_S;
    
    private byte[] I_C;
    
    /** Creates a new instance of DHKex */
    public DHKex() {
        this(1);
    }
    /**
     *  Init a new DHKex object with a specific operation code
     *  @param opCode
     *         1 - group1;
     */
    public DHKex(int opCode){
        switch(opCode){
            case 1:
                group1();
        }
    }
    private void group1(){
        do{
            setP(new BigInteger("ffffffffffffffffc90fdaa22168c234c4c6628b80dc1cd129024e088a67cc74020bbea63b139b22514a08798e3404ddef9519b3cd3a431b302b0a6df25f14374fe1356d6d51c245e485b576625e7ec6f44c42e9a637ed6b0bff5cb6f406b7edee386bfb5a899fa5ae9f24117c4b1fe649286651ece65381ffffffffffffffff", 16));
            setG(new BigInteger("2", 16));
            int qLength = Math.min(512, getP().bitLength() - 1 - 1);
            setX(new BigInteger(qLength, new Random()));
            setY(getG().modPow(getX(), getP()));
        }while(!useageCheck(getY(), getP()));
    }
    
    /** 回傳client的public key **/
    public synchronized byte[] getE(){
        if(e == null){
            e = getY().toByteArray();
        }
        return e;
    }
    
    /** 檢察產生出來的值是否可使用
     *  @see RFC4419
     */
    private boolean useageCheck(BigInteger test, BigInteger max){
        if(test.compareTo(BigInteger.ONE) == 1 && test.compareTo(max) == -1){
            return true;
        }else{
            return false;
        }
    }
    
    /** pass a new packet to this object **/
    public void newPacket(SshPacketV2 pn){
        this.setPn(pn);
    }
    
    /** 當SshIO收到server回傳 dhkex reply時執行 **/
    public boolean dhkex_reply() throws SSHConnectionException{
        if(getPn() == null){
            throw new SSHConnectionException("Connection terminated");
        }else{
            try {
                /** dhkex reply packet include:
                 *  String K_S
                 *  mping f
                 *  String sign_of_H
                 */
                pn.popByte(); // first byte (msg code)
                setK_S(pn.popBytes());  // dismount K_S
                
                byte[] f = pn.popBytes();  // dismount f
                
                setK(new BigInteger(f).modPow(getX(), getP()).toByteArray());
                
                setSign_of_H(pn.popBytes()); // dismount sign_of_H
                /** construct a psudo-SSHpacket **/
                
                //The hash H is computed as the HASH hash of the concatenation
                // of the
                //following:
                // string V_C, the client's version string (CR and NL excluded)
                // string V_S, the server's version string (CR and NL excluded)
                // string I_C, the payload of the client's SSH_MSG_KEXINIT
                // string I_S, the payload of the server's SSH_MSG_KEXINIT
                // string K_S, the host key
                // mpint e, exchange value sent by the client
                // mpint f, exchange value sent by the server
                // mpint K, the shared secret
                // This value is called the exchange hash, and it is used to
                // authenticate the key exchange.

                SshPacketV2 px = new SshPacketV2();
                px.pushString(getV_C());
                px.pushString(getV_S());
                px.pushString(getI_C());
                px.pushString(getI_S());
                px.pushString(getK_S());
                px.pushString(getE());
                px.pushString(f);
                px.pushString(getK());
                byte[] temp = px.getPayload();                
                
                SHA1 sha = new SHA1();
                sha.update(temp, 0, temp.length);
                setH(new byte[sha.getDigestSize()]);
                sha.doFinal(getH(), 0);
                
                //** dismount K_S **//
                /**      string    "ssh-dss"
                 * mpint     p
                 * mpint     q
                 * mpint     g
                 * mpint     y
                 ***/
                pn = new SshPacketV2();
                pn.setPayload(getK_S());
                
                String alg = pn.popString();
 
                if(alg.equals("ssh-dss")){
                    return verifyDSA(getH(), getSign_of_H(), new BigInteger(pn.popBytes()), new BigInteger(pn.popBytes()), new BigInteger(pn.popBytes()), new BigInteger(pn.popBytes()));
                }else{
                    throw new SSHConnectionException("Algorithm confused");
                }
                
            } catch (IllegalArgumentException ex) {
                ex.printStackTrace();
            } catch( OutOfMemoryError o){
                o.printStackTrace();
            }
            
        }
        return false;
    }
    
    
    /** reference: MidpSSH DHKeyExchange class
     *
     *  for more license information, please visit the MidpSSH software official website */
    private boolean verifyDSA(byte[] msg, byte[] sign, BigInteger p, BigInteger q, BigInteger g, BigInteger y) throws SSHConnectionException, NumberFormatException, OutOfMemoryError{
        /** sign format:
         *  String algorithm
         *  mpint blob
         */
        pn = new SshPacketV2();
        pn.setPayload(sign);
        String alg = pn.popString();   // alg name
        if( !alg.equals("ssh-dss")){
            throw new SSHConnectionException("Algorithm confused");
        }
        
        byte[] blob = pn.popBytes(); // dismount server blob
        /** end dismount **/
        
        int rslen = blob.length / 2;
        
        byte[] temp = new byte[rslen];
        temp[0] = 0;
        System.arraycopy(blob, 0, temp, 0, rslen);
        BigInteger r = new BigInteger(1, temp);
        System.arraycopy(blob, rslen, temp, 0, rslen);
        BigInteger s = new BigInteger(1, temp);
        
        /** verifySign(_msg, r, s)**/
        SHA1 sha1 = new SHA1();
        sha1.update(msg, 0, msg.length);
        byte[] hash = new byte[sha1.getDigestSize()];
        sha1.doFinal(hash, 0);
        msg = hash;
        
        BigInteger m = new BigInteger(1, msg);
        BigInteger zero = BigInteger.valueOf(0);
        
        if (zero.compareTo(r) >= 0 || q.compareTo(r) <= 0){
            
            return false;
        }
        if (zero.compareTo(s) >= 0 || q.compareTo(s) <= 0){
            return false;
        }
            
        BigInteger w = s.modInverse(q);
        
        BigInteger u1 = m.multiply(w).mod(q);
        BigInteger u2 = r.multiply(w).mod(q);
        
        u1 = g.modPow(u1, p);
        u2 = y.modPow(u2, p);
        
        BigInteger v = u1.multiply(u2).mod(p).mod(q);
        
        return v.equals(r);
    }
    
    
    /////////////////// Getter / Setter ////////////////////////////////
    ////////////////////////////////////////////////////////////////////

    public BigInteger getP() {
        return p;
    }

    public void setP(BigInteger p) {
        this.p = p;
    }

    public BigInteger getG() {
        return g;
    }

    public void setG(BigInteger g) {
        this.g = g;
    }

    public BigInteger getX() {
        return x;
    }

    public void setX(BigInteger x) {
        this.x = x;
    }

    public BigInteger getY() {
        return y;
    }

    public void setY(BigInteger y) {
        this.y = y;
    }

    public void setE(byte[] e) {
        this.e = e;
    }

    public byte[] getK_S() {
        return K_S;
    }

    public void setK_S(byte[] K_S) {
        this.K_S = K_S;
    }

    public byte[] getSign_of_H() {
        return sign_of_H;
    }

    public void setSign_of_H(byte[] sign_of_H) {
        this.sign_of_H = sign_of_H;
    }

    public byte[] getK() {
        return K;
    }

    public void setK(byte[] K) {
        this.K = K;
    }

    public byte[] getV_S() {
        return V_S;
    }

    public void setV_S(byte[] V_S) {
        this.V_S = V_S;
    }

    public byte[] getH() {
        return H;
    }

    public void setH(byte[] H) {
        this.H = H;
    }

    public byte[] getV_C() {
        return V_C;
    }

    public void setV_C(byte[] V_C) {
        this.V_C = V_C;
    }

    public byte[] getI_S() {
        return I_S;
    }

    public void setI_S(byte[] I_S) {
        this.I_S = I_S;
    }

    public byte[] getI_C() {
        return I_C;
    }

    public void setI_C(byte[] I_C) {
        this.I_C = I_C;
    }
    public void setPn(SshPacketV2 pn){
        this.pn = pn;
    }
    public SshPacketV2 getPn(){
        return pn;
    }
}
